import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import cv2 as cv
import numpy as np
from datetime import datetime

class RectangleSelector:
    def __init__(self, root):
        self.root = root
        self.root.title("图像矩形选择工具")
        self.root.geometry("1200x700")
        
        # 初始化变量
        self.current_img = None
        self.display_img = None
        self.photo = None
        self.rect_id = None
        self.start_x = None
        self.start_y = None
        self.selected_rect = None
        self.scale = 1.0
        self.offset_x = 0
        self.offset_y = 0
        
        # 保存选择结果的列表
        self.selection_results = []
        
        self.setup_ui()
        
    def setup_ui(self):
        # 创建菜单栏
        menubar = tk.Menu(self.root)
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="打开图像", command=self.open_image)
        file_menu.add_command(label="保存所有选择结果", command=self.save_all_selections)
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.root.quit)
        menubar.add_cascade(label="文件", menu=file_menu)
        self.root.config(menu=menubar)
        
        # 创建主框架
        main_frame = tk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 创建顶部框架（说明和主图像）
        top_frame = tk.Frame(main_frame)
        top_frame.pack(fill=tk.BOTH, expand=True)
        
        # 说明标签
        instruction_label = tk.Label(top_frame, text="1. 打开图像文件\n2. 在图像上拖拽鼠标绘制矩形选择区域\n3. 点击'应用选择'保存选择结果", 
                                   font=('Arial', 10), justify=tk.LEFT)
        instruction_label.pack(side=tk.TOP, pady=5)
        
        # 创建内容框架（包含主图像和结果预览）
        content_frame = tk.Frame(top_frame)
        content_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建画布框架
        canvas_frame = tk.Frame(content_frame)
        canvas_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        # 创建画布
        self.canvas = tk.Canvas(canvas_frame, bg='lightgray', cursor='cross')
        self.canvas.pack(fill=tk.BOTH, expand=True)
        
        # 创建右侧控制面板框架
        control_frame = tk.Frame(content_frame, width=250)
        control_frame.pack(side=tk.RIGHT, fill=tk.Y, padx=(10, 0))
        control_frame.pack_propagate(False)  # 固定宽度
        
        # 预览标签
        tk.Label(control_frame, text="选择区域预览:", font=('Arial', 9, 'bold')).pack(anchor=tk.W, pady=(0, 5))
        self.preview_label = tk.Label(control_frame, bg='white', width=20, height=10)
        self.preview_label.pack(pady=(0, 10))
        
        # 坐标信息标签
        tk.Label(control_frame, text="坐标信息:", font=('Arial', 9, 'bold')).pack(anchor=tk.W, pady=(0, 5))
        self.coord_label = tk.Label(control_frame, text="未选择区域", bg='lightyellow', wraplength=230)
        self.coord_label.pack(fill=tk.X, pady=(0, 10))
        
        # 按钮框架
        button_frame = tk.Frame(control_frame)
        button_frame.pack(fill=tk.X, pady=10)
        
        self.apply_button = tk.Button(button_frame, text="应用选择", state=tk.DISABLED, 
                                    command=self.apply_selection, width=12, bg='lightblue')
        self.apply_button.pack(side=tk.TOP, pady=2, fill=tk.X)
        
        self.reset_button = tk.Button(button_frame, text="重置选择", state=tk.DISABLED,
                                    command=self.reset_selection, width=12, bg='lightcoral')
        self.reset_button.pack(side=tk.TOP, pady=2, fill=tk.X)
        
        # 创建结果展示区域
        results_frame = tk.Frame(main_frame)
        results_frame.pack(fill=tk.BOTH, expand=False, pady=(10, 5), ipady=150)  # 固定高度
        
        # 结果标题
        results_label = tk.Label(results_frame, text="选择结果历史:", font=('Arial', 10, 'bold'))
        results_label.pack(fill=tk.X)
        
        # 创建结果展示容器（带滚动条）
        results_container = tk.Frame(results_frame)
        results_container.pack(fill=tk.BOTH, expand=True)
        
        # 创建画布和滚动条用于结果展示
        self.results_canvas = tk.Canvas(results_container, bg='white', height=120)  # 固定高度
        results_scrollbar = tk.Scrollbar(results_container, orient="horizontal", command=self.results_canvas.xview)
        self.results_canvas.configure(xscrollcommand=results_scrollbar.set)
        
        results_scrollbar.pack(side=tk.BOTTOM, fill=tk.X)
        self.results_canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
        
        # 创建结果展示框架
        self.results_inner_frame = tk.Frame(self.results_canvas)
        self.results_canvas.create_window((0, 0), window=self.results_inner_frame, anchor="nw")
        self.results_inner_frame.bind("<Configure>", self.on_results_frame_configure)
        
        # 绑定鼠标事件
        self.canvas.bind("<ButtonPress-1>", self.start_rect)
        self.canvas.bind("<B1-Motion>", self.draw_rect)
        self.canvas.bind("<ButtonRelease-1>", self.end_rect)
        
        # 状态栏
        self.status_bar = tk.Label(self.root, text="就绪", bd=1, relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
    def on_results_frame_configure(self, event):
        """更新滚动区域"""
        self.results_canvas.configure(scrollregion=self.results_canvas.bbox("all"))
        
    def open_image(self):
        file_path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp *.gif")]
        )
        
        if file_path:
            # 读取图像
            img = cv.imread(file_path)
            if img is not None:
                self.current_img = img
                self.display_image_on_canvas()
                self.status_bar.config(text=f"已打开图像: {file_path}")
            else:
                messagebox.showerror("错误", "无法读取图像文件")
                
    def display_image_on_canvas(self):
        if self.current_img is None:
            return
            
        # 获取画布尺寸
        canvas_width = self.canvas.winfo_width()
        canvas_height = self.canvas.winfo_height()
        
        # 如果画布尺寸为0，使用默认值
        if canvas_width <= 1 or canvas_height <= 1:
            canvas_width, canvas_height = 800, 500
            
        # 获取图像尺寸
        img_h, img_w = self.current_img.shape[:2]
        
        # 计算缩放比例
        self.scale = min(canvas_width/img_w, canvas_height/img_h) * 0.9  # 留一些边距
        display_w = int(img_w * self.scale)
        display_h = int(img_h * self.scale)
        
        # 缩放图像
        self.display_img = cv.resize(self.current_img, (display_w, display_h))
        
        # 转换为RGB格式
        if len(self.display_img.shape) == 3:
            rgb_img = cv.cvtColor(self.display_img, cv.COLOR_BGR2RGB)
        else:
            rgb_img = cv.cvtColor(self.display_img, cv.COLOR_GRAY2RGB)
            
        # 转换为PIL图像并在画布上显示
        pil_img = Image.fromarray(rgb_img)
        self.photo = ImageTk.PhotoImage(pil_img)
        
        # 清空画布
        self.canvas.delete("all")
        
        # 在画布中心显示图像
        self.offset_x = (canvas_width - display_w) // 2
        self.offset_y = (canvas_height - display_h) // 2
        self.canvas.create_image(self.offset_x, self.offset_y, anchor=tk.NW, image=self.photo)
        
        # 重置选择状态
        self.reset_selection()
        
    def start_rect(self, event):
        if self.current_img is None:
            return
            
        # 检查点击是否在图像区域内
        x, y = event.x, event.y
        img_right = self.offset_x + self.display_img.shape[1]
        img_bottom = self.offset_y + self.display_img.shape[0]
        
        if (self.offset_x <= x <= img_right and 
            self.offset_y <= y <= img_bottom):
            self.start_x = x
            self.start_y = y
            
            # 删除之前的矩形
            if self.rect_id:
                self.canvas.delete(self.rect_id)
                
            # 创建新矩形
            self.rect_id = self.canvas.create_rectangle(
                x, y, x, y,
                outline='red', width=2, dash=(4, 2)
            )
            
    def draw_rect(self, event):
        if self.rect_id and self.start_x is not None:
            self.canvas.coords(
                self.rect_id,
                self.start_x,
                self.start_y,
                event.x,
                event.y
            )
            
    def end_rect(self, event):
        if self.rect_id and self.start_x is not None:
            x1 = min(self.start_x, event.x)
            y1 = min(self.start_y, event.y)
            x2 = max(self.start_x, event.x)
            y2 = max(self.start_y, event.y)
            
            # 检查是否在图像区域内
            img_right = self.offset_x + self.display_img.shape[1]
            img_bottom = self.offset_y + self.display_img.shape[0]
            
            if (self.offset_x <= x1 <= img_right and 
                self.offset_y <= y1 <= img_bottom and
                self.offset_x <= x2 <= img_right and
                self.offset_y <= y2 <= img_bottom):
                
                # 更新矩形坐标
                self.canvas.coords(self.rect_id, x1, y1, x2, y2)
                self.selected_rect = (x1, y1, x2, y2)
                
                # 启用按钮
                self.apply_button.config(state=tk.NORMAL)
                self.reset_button.config(state=tk.NORMAL)
                
                # 显示预览和坐标
                self.show_preview(x1, y1, x2, y2)
                self.update_coordinates(x1, y1, x2, y2)
            else:
                # 删除无效矩形
                self.canvas.delete(self.rect_id)
                self.rect_id = None
                self.apply_button.config(state=tk.DISABLED)
                self.reset_button.config(state=tk.DISABLED)
                self.coord_label.config(text="未选择区域")
                self.preview_label.config(image='')
                
    def show_preview(self, x1, y1, x2, y2):
        if self.current_img is None:
            return
            
        # 计算在原始图像中的坐标
        orig_x1 = int((x1 - self.offset_x) / self.scale)
        orig_y1 = int((y1 - self.offset_y) / self.scale)
        orig_x2 = int((x2 - self.offset_x) / self.scale)
        orig_y2 = int((y2 - self.offset_y) / self.scale)
        
        # 确保坐标在有效范围内
        orig_x1 = max(0, orig_x1)
        orig_y1 = max(0, orig_y1)
        orig_x2 = min(self.current_img.shape[1], orig_x2)
        orig_y2 = min(self.current_img.shape[0], orig_y2)
        
        # 截取图像区域
        if orig_x2 > orig_x1 and orig_y2 > orig_y1:
            cropped_img = self.current_img[orig_y1:orig_y2, orig_x1:orig_x2]
            
            # 创建预览图像
            preview_h, preview_w = cropped_img.shape[:2]
            max_preview_size = 150
            preview_scale = min(max_preview_size/max(preview_w, preview_h), 1.0)
            preview_w = int(preview_w * preview_scale)
            preview_h = int(preview_h * preview_scale)
            
            if preview_w > 0 and preview_h > 0:
                preview_img = cv.resize(cropped_img, (preview_w, preview_h))
                
                # 转换为RGB格式
                if len(preview_img.shape) == 3:
                    preview_rgb = cv.cvtColor(preview_img, cv.COLOR_BGR2RGB)
                else:
                    preview_rgb = cv.cvtColor(preview_img, cv.COLOR_GRAY2RGB)
                
                # 转换为PIL图像并显示
                preview_pil = Image.fromarray(preview_rgb)
                preview_photo = ImageTk.PhotoImage(preview_pil)
                self.preview_label.config(image=preview_photo)
                self.preview_label.image = preview_photo  # 保持引用
                
    def update_coordinates(self, x1, y1, x2, y2):
        if self.current_img is None:
            return
            
        # 计算在原始图像中的坐标
        orig_x1 = int((x1 - self.offset_x) / self.scale)
        orig_y1 = int((y1 - self.offset_y) / self.scale)
        orig_x2 = int((x2 - self.offset_x) / self.scale)
        orig_y2 = int((y2 - self.offset_y) / self.scale)
        
        # 确保坐标在有效范围内
        orig_x1 = max(0, orig_x1)
        orig_y1 = max(0, orig_y1)
        orig_x2 = min(self.current_img.shape[1], orig_x2)
        orig_y2 = min(self.current_img.shape[0], orig_y2)
        
        coord_text = f"({orig_x1}, {orig_y1}) to ({orig_x2}, {orig_y2})"
        self.coord_label.config(text=coord_text)
        
    def apply_selection(self):
        if self.selected_rect and self.current_img is not None:
            x1, y1, x2, y2 = self.selected_rect
            
            # 计算在原始图像中的坐标
            orig_x1 = int((x1 - self.offset_x) / self.scale)
            orig_y1 = int((y1 - self.offset_y) / self.scale)
            orig_x2 = int((x2 - self.offset_x) / self.scale)
            orig_y2 = int((y2 - self.offset_y) / self.scale)
            
            # 确保坐标在有效范围内
            orig_x1 = max(0, orig_x1)
            orig_y1 = max(0, orig_y1)
            orig_x2 = min(self.current_img.shape[1], orig_x2)
            orig_y2 = min(self.current_img.shape[0], orig_y2)
            
            # 截取图像区域
            if orig_x2 > orig_x1 and orig_y2 > orig_y1:
                cropped_img = self.current_img[orig_y1:orig_y2, orig_x1:orig_x2]
                
                # 保存选择结果
                timestamp = datetime.now().strftime("%H:%M:%S")
                selection_info = {
                    'image': cropped_img.copy(),
                    'coordinates': (orig_x1, orig_y1, orig_x2, orig_y2),
                    'timestamp': timestamp
                }
                self.selection_results.append(selection_info)
                
                # 在结果展示区域显示
                self.display_selection_result(selection_info, len(self.selection_results)-1)
                
                # 更新状态栏
                self.status_bar.config(text=f"已保存选择结果 {len(self.selection_results)}: ({orig_x1}, {orig_y1}) to ({orig_x2}, {orig_y2})")
                
                # 重置选择
                self.reset_selection()
                
    def display_selection_result(self, selection_info, index):
        """在结果展示区域显示选择结果"""
        # 创建结果框架
        result_frame = tk.Frame(self.results_inner_frame, relief=tk.RAISED, borderwidth=1)
        result_frame.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 调整图像大小
        cropped_img = selection_info['image']
        img_h, img_w = cropped_img.shape[:2]
        max_size = 80  # 减小缩略图大小
        scale = min(max_size/max(img_w, img_h), 1.0)
        display_w = int(img_w * scale)
        display_h = int(img_h * scale)
        
        if display_w > 0 and display_h > 0:
            display_img = cv.resize(cropped_img, (display_w, display_h))
            
            # 转换为RGB格式
            if len(display_img.shape) == 3:
                rgb_img = cv.cvtColor(display_img, cv.COLOR_BGR2RGB)
            else:
                rgb_img = cv.cvtColor(display_img, cv.COLOR_GRAY2RGB)
                
            # 转换为PIL图像并显示
            pil_img = Image.fromarray(rgb_img)
            photo = ImageTk.PhotoImage(pil_img)
            
            img_label = tk.Label(result_frame, image=photo, bg='white')
            img_label.pack(padx=2, pady=2)
            img_label.image = photo  # 保持引用
            
        # 创建标题
        title = f"选择 {index+1}"
        title_label = tk.Label(result_frame, text=title, font=('Arial', 8))
        title_label.pack(padx=2, pady=(0, 2))
        
        # 创建按钮框架
        button_frame = tk.Frame(result_frame)
        button_frame.pack(fill=tk.X, padx=2, pady=(0, 2))
        # 查看按钮
        view_button = tk.Button(button_frame, text="查看",
                                command=lambda: self.view_selection_detail(selection_info),
                                width=4, height=1, font=('Arial', 7))
        view_button.pack(side=tk.LEFT, padx=1)

        # 删除按钮
        delete_button = tk.Button(button_frame, text="删",
                                  command=lambda: self.delete_selection(index, result_frame),
                                  width=4, height=1, font=('Arial', 7), bg='lightcoral')
        delete_button.pack(side=tk.LEFT, padx=1)

        # 添加匹配按钮
        match_button = tk.Button(button_frame, text="匹配",
                                 command=lambda: self.match_selection(selection_info['image']), width=6,
                                 bg='lightgreen')
        match_button.pack(side=tk.LEFT, padx=2)
        # # 匹配按钮
        # match_button = tk.Button(button_frame, text="匹配",
        #                          command=lambda: self.match_selection(selection_info),
        #                          width=4, height=1, font=('Arial', 7), bg='lightcoral')
        # match_button.pack(side=tk.LEFT, padx=1)

        # 更新滚动区域
        self.results_inner_frame.update_idletasks()
        self.results_canvas.configure(scrollregion=self.results_canvas.bbox("all"))


        
    def match_selection(self, template_img):
        """将选择的图像匹配到原图中"""
        if self.current_img is None or template_img is None:
            messagebox.showwarning("警告", "请先打开图像并选择区域")
            return
        
        # 确保模板图像和原图类型一致
        if len(template_img.shape) != len(self.current_img.shape):
            messagebox.showerror("错误", "图像类型不匹配")
            return
        
        # 使用模板匹配
        result = cv.matchTemplate(self.current_img, template_img, cv.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv.minMaxLoc(result)
        
        # 获取匹配区域的坐标
        top_left = max_loc
        h, w = template_img.shape[:2]
        bottom_right = (top_left[0] + w, top_left[1] + h)
        
        # 在画布上高亮显示匹配区域
        self.canvas.delete("match_rect")
        self.canvas.create_rectangle(
            top_left[0] * self.scale + self.offset_x,
            top_left[1] * self.scale + self.offset_y,
            bottom_right[0] * self.scale + self.offset_x,
            bottom_right[1] * self.scale + self.offset_y,
            outline='blue', width=2, tags="match_rect"
        )
        
        # 更新状态栏
        self.status_bar.config(text=f"匹配成功: 位置 ({top_left[0]}, {top_left[1]}) 到 ({bottom_right[0]}, {bottom_right[1]})")
        


    def view_selection_detail(self, selection_info):
        """查看选择结果的详细信息"""
        # 创建新窗口
        detail_window = tk.Toplevel(self.root)
        detail_window.title(f"选择结果详情")
        detail_window.geometry("600x500")
        
        # 居中显示新窗口
        detail_window.update_idletasks()
        width = detail_window.winfo_width()
        height = detail_window.winfo_height()
        x = (detail_window.winfo_screenwidth() // 2) - (width // 2)
        y = (detail_window.winfo_screenheight() // 2) - (height // 2)
        detail_window.geometry('{}x{}+{}+{}'.format(width, height, x, y))
        
        # 创建主框架
        main_frame = tk.Frame(detail_window)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 创建信息标签
        coords = selection_info['coordinates']
        info_text = f"选择区域: ({coords[0]}, {coords[1]}) to ({coords[2]}, {coords[3]})\n时间: {selection_info['timestamp']}"
        info_label = tk.Label(main_frame, text=info_text, font=('Arial', 10, 'bold'))
        info_label.pack(pady=5)
        
        # 创建画布框架
        canvas_frame = tk.Frame(main_frame)
        canvas_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # 创建画布用于显示结果图像
        detail_canvas = tk.Canvas(canvas_frame, bg='lightgray')
        detail_canvas.pack(fill=tk.BOTH, expand=True)
        
        # 调整图像大小以适应画布
        cropped_img = selection_info['image']
        img_h, img_w = cropped_img.shape[:2]
        canvas_width = 500
        canvas_height = 400
        
        # 计算缩放比例
        scale = min(canvas_width/img_w, canvas_height/img_h) * 0.9
        display_w = int(img_w * scale)
        display_h = int(img_h * scale)
        
        # 缩放图像
        display_img = cv.resize(cropped_img, (display_w, display_h))
        
        # 转换为RGB格式
        if len(display_img.shape) == 3:
            rgb_img = cv.cvtColor(display_img, cv.COLOR_BGR2RGB)
        else:
            rgb_img = cv.cvtColor(display_img, cv.COLOR_GRAY2RGB)
            
        # 转换为PIL图像并在画布上显示
        pil_img = Image.fromarray(rgb_img)
        photo = ImageTk.PhotoImage(pil_img)
        
        # 在画布中心显示图像
        img_x = (canvas_width - display_w) // 2
        img_y = (canvas_height - display_h) // 2
        detail_canvas.create_image(img_x, img_y, anchor=tk.NW, image=photo)
        
        # 保持图像引用
        detail_canvas.image = photo
        
        # 添加保存按钮
        button_frame = tk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=10)
        
        save_button = tk.Button(button_frame, text="保存图像", 
                               command=lambda: self.save_single_selection(selection_info),
                               width=12, bg='lightgreen')
        save_button.pack(side=tk.LEFT, padx=10)
        
        close_button = tk.Button(button_frame, text="关闭", 
                                command=detail_window.destroy,
                                width=12, bg='lightcoral')
        close_button.pack(side=tk.RIGHT, padx=10)
        
    def save_single_selection(self, selection_info):
        """保存单个选择结果"""
        file_path = filedialog.asksaveasfilename(
            defaultextension=".png",
            filetypes=[("PNG图像", "*.png"), ("JPEG图像", "*.jpg"), ("BMP图像", "*.bmp"), ("所有文件", "*.*")]
        )
        
        if file_path:
            try:
                cv.imwrite(file_path, selection_info['image'])
                messagebox.showinfo("成功", f"图像已保存到: {file_path}")
            except Exception as e:
                messagebox.showerror("错误", f"保存图像失败: {str(e)}")
                
    def save_all_selections(self):
        """保存所有选择结果"""
        if not self.selection_results:
            messagebox.showinfo("提示", "没有选择结果可保存")
            return
            
        # 选择保存目录
        directory = filedialog.askdirectory()
        if directory:
            try:
                saved_count = 0
                for i, selection_info in enumerate(self.selection_results):
                    filename = f"selection_{i+1}_{selection_info['timestamp'].replace(':', '-')}.png"
                    file_path = f"{directory}/{filename}"
                    cv.imwrite(file_path, selection_info['image'])
                    saved_count += 1
                    
                messagebox.showinfo("成功", f"已保存 {saved_count} 个选择结果到: {directory}")
            except Exception as e:
                messagebox.showerror("错误", f"保存图像失败: {str(e)}")
                
    def delete_selection(self, index, frame):
        """删除选择结果"""
        if 0 <= index < len(self.selection_results):
            # 从列表中删除
            del self.selection_results[index]
            
            # 销毁框架
            frame.destroy()
            
            # 更新状态栏
            self.status_bar.config(text=f"已删除选择结果 {index+1}，剩余 {len(self.selection_results)} 个结果")
            
            # 更新滚动区域
            self.results_inner_frame.update_idletasks()
            self.results_canvas.configure(scrollregion=self.results_canvas.bbox("all"))
            
    def reset_selection(self):
        # 删除矩形
        if self.rect_id:
            self.canvas.delete(self.rect_id)
            self.rect_id = None
            
        # 重置变量
        self.start_x = None
        self.start_y = None
        self.selected_rect = None
        
        # 禁用按钮
        self.apply_button.config(state=tk.DISABLED)
        self.reset_button.config(state=tk.DISABLED)
        
        # 清空预览和坐标信息
        self.preview_label.config(image='')
        self.coord_label.config(text="未选择区域")

def main():
    root = tk.Tk()
    app = RectangleSelector(root)
    root.mainloop()

if __name__ == "__main__":
    main()